package com.btl.service.ai.support.input.common.split.splitters;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.segment.TextSegment;

import java.util.ArrayList;
import java.util.List;
import java.util.regex.Pattern;

public class FaqQASplitter implements DocumentSplitter {

    private static final String QUESTION_MARK = "Q:";
    private static final String ANSWER_MARK = "A:";

    private final String questionMark;
    private final String answerMark;

    public FaqQASplitter() {
        this(QUESTION_MARK);
    }

    public FaqQASplitter(String questionMark) {
        this.questionMark = questionMark;
        this.answerMark = ANSWER_MARK;
    }

    @Override
    public List<TextSegment> split(Document document) {
        String text = document.text();
        String qRegex = "(?=" + Pattern.quote(questionMark) + ")";
        String[] parts = text.split(qRegex);
        List<TextSegment> segments = new ArrayList<>();
        for (String part : parts) {
            String trimmed = part.trim();
            if (!trimmed.isEmpty()) {
                String aRegex = "(?=" + Pattern.quote(answerMark) + ")";
                String[] subParts = trimmed.split(aRegex);
                Metadata metadata = document.metadata();
                if (subParts.length == 2) {
                    metadata = metadata.copy()
                            .put("title", subParts[0]);
                    segments.add(TextSegment.from(subParts[1], metadata));
                } else {
                    segments.add(TextSegment.from(trimmed, metadata));
                }
            }
        }
        return segments;
    }
}
